# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import collections
import hashlib
import json
import logging
import threading
from contextlib import contextmanager

import native.prophet.py.prophet.graph as native_graph
from experimental.prophet import op_registry
from experimental.prophet.graph import (graph_def_pb2, variable as var_mod)
from experimental.prophet.graph.shape import Shape
from experimental.prophet.graph.util import (from_data_type_proto,
                                             to_data_type_proto,
                                             build_output_tuple)


def as_variable(v):
  return var_mod.as_variable(v)


class GraphBuildingContext:
  def __init__(self, graph):
    self._graph = graph
    self._namespace_stack = ''
    self._control_input = None
    self._default_timestamp = None
    self._aux_dict = {}

  @property
  def aux_dict(self):
    return self._aux_dict

  def create_node(self, op, inputs, attrs=None, exec_attrs=None, name=None, namespace=None):
    name = self._get_unique_node_name(name or op, namespace=namespace)
    exec_attrs = exec_attrs if exec_attrs is not None else {}
    return self._graph.create_node(op=op,
                                   inputs=inputs,
                                   control_input=self.control_input,
                                   attrs=attrs,
                                   exec_attrs=exec_attrs,
                                   name=name)

  @contextmanager
  def namespace(self, name):
    assert name and '/' not in name
    old_stack = self._namespace_stack
    try:
      self._namespace_stack = (old_stack + '/' + name) if old_stack else name
      yield self._namespace_stack
    finally:
      self._namespace_stack = old_stack

  def get_namespace(self):
    return self._namespace_stack

  def _get_unique_node_name(self, name, namespace=None):
    assert name
    if namespace is None:
      namespace = self.get_namespace()
    if namespace:
      name = namespace + '/' + name
    base_name = name
    i = 1
    while name in self._graph.node_map:
      i += 1
      name = '%s_%d' % (base_name, i)
    return name

  def register_named_variable_ref(self, var_ref):
    if var_ref.name in self._graph.named_vars:
      if var_ref.unref() == self._graph.named_vars[var_ref.name].unref():
        return
      logging.warning('Variable name `%s` already exists.' % var_ref.name)
    self._graph.named_vars[var_ref.name] = var_ref

  @property
  def control_input(self):
    return self._control_input

  def _control_input_and(self, pred1, pred2):
    if pred1 is None:
      return pred2
    assert pred2 is not None

    name = self._get_unique_node_name('control_if')
    node = self._graph.create_node(op='ElemwiseBinOp',
                                   inputs=[pred1, pred2],
                                   control_input=None,
                                   name=name,
                                   exec_attrs={},
                                   attrs={
                                       'dtype1': graph_def_pb2.DATA_TYPE_BOOL,
                                       'dtype2': graph_def_pb2.DATA_TYPE_BOOL,
                                       'func': 'And'
                                   })
    return node.outputs[0]

  @contextmanager
  def control_if(self, pred):
    if pred is not None:
      pred = as_variable(pred).cast(bool)
      assert Shape(pred.shape).is_scalar

    old_control_input = self._control_input
    try:
      if pred is not None:
        self._control_input = self._control_input_and(old_control_input, pred)
      else:
        self._control_input = None
      yield
    finally:
      self._control_input = old_control_input

  @property
  def default_timestamp(self):
    if self._default_timestamp is None:
      from experimental.prophet.ops.globals import global_timestamp
      self._default_timestamp = global_timestamp()
    return self._default_timestamp

  @contextmanager
  def as_default_timestamp(self, timestamp):
    assert timestamp is not None
    old_timestamp = self.default_timestamp
    try:
      self._default_timestamp = timestamp
      yield
    finally:
      self._default_timestamp = old_timestamp


_default_graph = threading.local()


def get_default_graph():
  try:
    graph = _default_graph.graph
    if graph is not None:
      return graph
  except AttributeError:
    pass
  graph = Graph()
  _default_graph.graph = graph
  return graph


def reset_default_graph():
  _default_graph.graph = None


@contextmanager
def as_default_graph(graph):
  old_graph = get_default_graph()
  try:
    _default_graph.graph = graph
    yield graph
  finally:
    _default_graph.graph = old_graph


def namespace(name):
  return get_default_graph().building_ctx.namespace(name)


def control_if(pred):
  return get_default_graph().building_ctx.control_if(pred)


def create_node(*args, **kwargs):
  return get_default_graph().building_ctx.create_node(*args, **kwargs)


def get_named_variable(name):
  return get_default_graph().get_named_variable(name)


def default_timestamp():
  return get_default_graph().building_ctx.default_timestamp


def as_default_timestamp(timestamp):
  return get_default_graph().building_ctx.as_default_timestamp(timestamp)


# Internal
def get_aux_entry(entry, graph=None, type_cls=dict):
  if graph is None:
    graph = get_default_graph()
  graph_aux_dict = graph.building_ctx.aux_dict

  if entry not in graph_aux_dict:
    graph_aux_dict[entry] = type_cls()
  return graph_aux_dict[entry]


def _shuffle_inputs(inputs, input_def, node_name):
  if isinstance(inputs, collections.Mapping):
    if input_def.variadic_input:
      new_inputs = [(arg_name, input.unref()) for arg_name, input in inputs.items()]
    else:
      assert set(inputs.keys()) <= set([arg.name for arg in input_def.input_args])
      new_inputs = []
      for arg in input_def.input_args:
        input = inputs[arg.name].unref()
        if to_data_type_proto(input.dtype) != arg.dtype:
          logging.warning('dtype mismatch on %s:%s and %s. Auto-typecasting',
                          node_name,
                          arg_name,
                          inputs[idx].instance_name)
          input = input.cast(from_data_type_proto(arg.dtype))
        new_inputs.append((arg.name, input))
      return new_inputs

  else:
    if input_def.variadic_input:
      new_inputs = [(str(idx), input.unref()) for idx, input in enumerate(inputs)]
    else:
      assert len(inputs) == len(input_def.input_args), ('%d!=%d' %
                                                        (len(inputs), len(input_def.input_args)))
      new_inputs = []
      for idx, arg in enumerate(input_def.input_args):
        arg_name = arg.name or str(idx)
        input = inputs[idx].unref()
        if to_data_type_proto(input.dtype) != arg.dtype:
          logging.warning('dtype mismatch on %s:%s and %s. Auto-typecasting',
                          node_name,
                          arg_name,
                          inputs[idx].instance_name)
          input = input.cast(from_data_type_proto(arg.dtype))
        new_inputs.append((arg_name, input))

  return collections.OrderedDict(new_inputs)


class Graph:
  def __init__(self):
    self.nodes = []
    self.node_map = {}
    self.named_vars = {}
    self._building_ctx = None
    self._instance = None

  @property
  def building_ctx(self):
    if self._building_ctx is None:
      self._building_ctx = GraphBuildingContext(self)
    return self._building_ctx

  def get_node(self, name):
    return self.node_map[name]

  def get_named_variable(self, name):
    return self.named_vars[name]

  def create_node(self, op, inputs, control_input, attrs, exec_attrs, name):
    assert name not in self.node_map
    input_def = op_registry.get_input_args(op, attrs, node_name=name)
    inputs = _shuffle_inputs(inputs, input_def, node_name=name)
    output_def = op_registry.get_output_args(op, attrs, inputs, node_name=name)

    op_def = op_registry.get_op_def(op)
    assert op_def.name == op
    op_def.variadic_input = input_def.variadic_input
    op_def.input_args.extend(input_def.input_args)
    op_def.output_args.extend(output_def.output_args)

    if control_input is not None:
      control_input = control_input.unref()

    node = Node(name=name,
                graph=self,
                inputs=inputs,
                control_input=control_input,
                op_def=op_def,
                attrs=attrs,
                exec_attrs=exec_attrs)
    self.nodes.append(node)
    self.node_map[name] = node
    return node

  def dump_dot(self):
    # Dump graph in dot language: https://www.graphviz.org
    lines = []
    lines.append('digraph {')
    lines.append('  %s;' % ' '.join([('"%s"' % node.name) for node in self.nodes]))
    for to_node in self.nodes:
      for input_arg, input_var in to_node.inputs.items():
        lines.append('  "%s" -> "%s" [ label = "%s" ];' %
                     (input_var.node.name, to_node.name, input_var.name or input_var.arg_name))
    lines.append('}')
    return '\n'.join(lines)

  def get_graph_def(self):
    graph_def = graph_def_pb2.GraphDefProto()
    graph_def.nodes.extend([node.get_node_def() for node in self.nodes])
    return graph_def

  def get_graph_hash(self):
    graph_def = self.get_graph_def()
    exec_attrs = [node.exec_attrs for node in self.nodes]
    return hashlib.sha256(graph_def.SerializeToString()).hexdigest()

  def eval(self):
    if not self.is_instantiated:
      self.instantiate()
    self._instance.eval()

  def instantiate(self):
    assert not self.is_instantiated
    s = self.get_graph_def().SerializeToString()
    self._instance = native_graph._PyGraph(s)

  @property
  def is_instantiated(self):
    return self._instance is not None

  def reset_instance(self):
    for node in self.nodes:
      for out_var in node.outputs:
        out_var._reset_instance()
    self._instance = None

  def _get_variable_instance(self, name):
    return self._instance.get_variable(name)


class Node:
  def __init__(self, graph, inputs, control_input, op_def, attrs, exec_attrs, name):
    self.graph = graph
    self.name = name

    self.inputs = inputs if inputs is not None else []
    self._outputs = build_output_tuple(
        [var_mod.Variable(self, idx, arg_def) for idx, arg_def in enumerate(op_def.output_args)])
    self.control_input = control_input

    self.op_def = op_def
    self.attrs = attrs
    self.exec_attrs = exec_attrs

  @property
  def outputs(self):
    return self._outputs

  @property
  def op(self):
    return self.op_def.name

  @property
  def is_constant(self):
    return self.op == 'Constant'

  @property
  def is_placeholder(self):
    return self.op == 'Placeholder'

  @property
  def is_shared_variable(self):
    return self.op == 'SharedVariable'

  def get_node_def(self):
    node_def = graph_def_pb2.NodeDefProto()
    node_def.name = self.name
    node_def.op = self.op_def.name
    node_def.op_def.CopyFrom(self.op_def)

    for input_name, input_var in self.inputs.items():
      arg_ref = node_def.inputs.add()
      arg_ref.node_name = input_var.node.name
      arg_ref.output_index = input_var.output_idx
      arg_ref.input_name = input_name
      arg_ref.dtype = to_data_type_proto(input_var.dtype)
      arg_ref.shape.CopyFrom(Shape(input_var.shape).to_proto())

    if self.control_input:
      arg_ref = graph_def_pb2.ArgRefProto()
      arg_ref.node_name = self.control_input.node.name
      arg_ref.output_index = self.control_input.output_idx
      arg_ref.dtype = to_data_type_proto(self.control_input.dtype)
      arg_ref.shape.CopyFrom(Shape(self.control_input.shape).to_proto())
      node_def.control_input.CopyFrom(arg_ref)

    if self.attrs is not None:
      node_def.attrs_json = json.dumps(self.attrs)
    return node_def
